import os
import re
from tqdm import tqdm

def delete_augmented_images(folder_path):
    """
    删除指定文件夹中文件名包含"aug"的图片文件（.jpg和.jpeg）
    
    参数:
    folder_path (str): 要处理的文件夹路径
    """
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"错误: 文件夹 '{folder_path}' 不存在")
        return
    
    # 定义正则表达式模式，匹配包含"aug"的.jpg和.jpeg文件
    pattern = re.compile(r'aug.*\.(jpg|jpeg)$', re.IGNORECASE)
    
    # 获取所有匹配的文件
    files_to_delete = []
    for filename in os.listdir(folder_path):
        if pattern.search(filename):
            file_path = os.path.join(folder_path, filename)
            if os.path.isfile(file_path):
                files_to_delete.append(file_path)
    
    # 确认是否要删除
    if not files_to_delete:
        print("没有找到需要删除的文件")
        return
    
    print(f"找到 {len(files_to_delete)} 个需要删除的文件")
    confirm = input(f"确定要删除这些文件吗？这将无法恢复。(y/N): ")
    
    if confirm.lower() != 'y':
        print("操作已取消")
        return
    
    # 使用tqdm显示进度条并删除文件
    for file_path in tqdm(files_to_delete, desc="删除进度", unit="文件"):
        try:
            os.remove(file_path)
        except Exception as e:
            print(f"无法删除文件 '{file_path}': {e}")
    
    print(f"操作完成，已删除 {len(files_to_delete)} 个文件")

if __name__ == "__main__":
    # 直接在代码中指定要处理的文件夹路径
    folder_path = "dataset_images_gongzhuang/train"  # 请替换为实际的文件夹路径
    
    delete_augmented_images(folder_path)    